import torch.nn as nn
import torch
import torch.nn.functional as F
from models.component.ContextualLoss import ContextualLoss_forward
from models.component.loss import GANLoss
from models.component.vgg19 import VGG19_feature_color_torchversion
from models.network.correspondence import NoVGGCorrespondence
from models.network.discriminator import MultiscaleDiscriminator
from models.network.generator import SPADEGenerator, DomainClassifier
from utils import util


class Backbone(nn.Module):
    def __init__(self, opt):
        super(Backbone, self).__init__()
        self.opt = opt
        net = self.initialize_networks()
        self.net = torch.nn.ModuleDict(net)
        self.FloatTensor = torch.cuda.FloatTensor if opt.device == 'cuda' else torch.FloatTensor

        if opt.phase == 'train':
            self.vggnet_fix = VGG19_feature_color_torchversion(vgg_normal_correct=opt.vgg_normal_correct)
            self.vggnet_fix.load_state_dict(torch.load('models/vgg19_conv.pth'))
            self.vggnet_fix.eval()
            for param in self.vggnet_fix.parameters():
                param.requires_grad = False
            self.vggnet_fix.to(opt.device)
            self.contextual_forward_loss = ContextualLoss_forward(opt)
            self.criterionGAN = GANLoss(
                opt.gan_mode, tensor=self.FloatTensor, opt=opt)
            self.criterionFeat = torch.nn.L1Loss()
            self.MSE_loss = torch.nn.MSELoss()
            if opt.which_perceptual == '5_2':
                self.perceptual_layer = -1
            elif opt.which_perceptual == '4_2':
                self.perceptual_layer = -2

    def save(self, epoch):
        util.save_network(self.net['netG'], 'G', epoch, self.opt)
        util.save_network(self.net['netDs'], 'Ds', epoch, self.opt)
        util.save_network(self.net['netDc'], 'Dc', epoch, self.opt)
        util.save_network(self.net['netCorr'], 'Corr', epoch, self.opt)
        if self.opt.weight_domainC > 0 and self.opt.domain_rela:
            util.save_network(self.net['netDomainClassifier'], 'DomainClassifier', epoch, self.opt)

    def forward(self, data, mode, GforD=None, alpha=1):
        input_label, input_semantics, real_image, self_ref, ref_image, ref_label, ref_semantics, similar_ref = self.preprocess_input(
            data)
        self.alpha = alpha
        generated_out = {}
        if mode == 'generator':
            g_loss, generated_out = self.compute_generator_loss(input_label,
                                                                input_semantics, real_image, ref_label, ref_semantics,
                                                                ref_image, self_ref, similar_ref)
            out = {}
            out['fake_image'] = generated_out['fake_image']
            out['input_semantics'] = input_semantics
            out['ref_semantics'] = ref_semantics
            out['warp_out'] = None if 'warp_out' not in generated_out else generated_out['warp_out']
            out['warp_mask'] = None if 'warp_mask' not in generated_out else generated_out['warp_mask']
            out['adaptive_feature_seg'] = None if 'adaptive_feature_seg' not in generated_out else generated_out[
                'adaptive_feature_seg']
            out['adaptive_feature_img'] = None if 'adaptive_feature_img' not in generated_out else generated_out[
                'adaptive_feature_img']
            out['warp_cycle'] = None if 'warp_cycle' not in generated_out else generated_out['warp_cycle']
            out['warp_i2r'] = None if 'warp_i2r' not in generated_out else generated_out['warp_i2r']
            out['warp_i2r2i'] = None if 'warp_i2r2i' not in generated_out else generated_out['warp_i2r2i']
            return g_loss, out

        if mode == 'discriminator':
            self.opt.D_type = 'structure'
            d_loss_structure = self.compute_discriminator_loss(
                input_semantics, real_image, GforD, label=input_label)
            self.opt.D_type = 'color'
            d_loss_color = self.compute_discriminator_loss(
                similar_ref, ref_image, GforD, label=input_label)

            D_losses = {}
            D_losses['D_Fake'] = d_loss_structure['D_Fake'] + d_loss_color['D_Fake']
            D_losses['D_real'] = d_loss_structure['D_real'] + d_loss_color['D_real']

            return D_losses
        elif mode == 'inference':
            out = {}
            with torch.no_grad():
                out = self.inference(input_semantics,
                                     ref_semantics=ref_semantics, ref_image=ref_image, self_ref=self_ref)
            out['input_semantics'] = input_semantics
            out['ref_semantics'] = ref_semantics
            return out
        else:
            raise ValueError("|mode| is invalid")

    def inference(self, input_semantics, ref_semantics=None, ref_image=None, self_ref=None):
        generate_out = {}
        coor_out = self.net['netCorr'](ref_image, None, input_semantics, ref_semantics)
        if self.opt.CBN_intype == 'mask':
            CBN_in = input_semantics
        elif self.opt.CBN_intype == 'warp':
            CBN_in = coor_out['warp_out']
        elif self.opt.CBN_intype == 'warp_mask':
            CBN_in = torch.cat((coor_out['warp_out'], input_semantics), dim=1)

        generate_out['fake_image'] = self.net['netG'](input_semantics, warp_out=CBN_in)
        generate_out = {**generate_out, **coor_out}
        return generate_out

    def compute_discriminator_loss(self, input_semantics, real_image, GforD, label=None):
        D_losses = {}
        with torch.no_grad():
            # fake_image, _, _, _, _ = self.generate_fake(input_semantics, real_image, VGG_feat=False)
            fake_image = GforD['fake_image'].detach()
            fake_image.requires_grad_()

        pred_fake, pred_real, seg, fake_cam_logit, real_cam_logit = self.discriminate(
            input_semantics, fake_image, real_image)

        D_losses['D_Fake'] = self.criterionGAN(pred_fake, False,
                                               for_discriminator=True) * self.opt.weight_gan
        D_losses['D_real'] = self.criterionGAN(pred_real, True,
                                               for_discriminator=True) * self.opt.weight_gan

        return D_losses

    def encode_z(self, real_image):
        mu, logvar = self.net['netE'](real_image)
        z = self.reparameterize(mu, logvar)
        return z, mu, logvar

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return eps.mul(std) + mu

    # Take the prediction of fake and real images from the combined batch
    def divide_pred(self, pred):
        # the prediction contains the intermediate outputs of multiscale GAN,
        # so it's usually a list
        if type(pred) == list:
            fake = []
            real = []
            for p in pred:
                fake.append([tensor[:tensor.size(0) // 2] for tensor in p])
                real.append([tensor[tensor.size(0) // 2:] for tensor in p])
        else:
            fake = pred[:pred.size(0) // 2]
            real = pred[pred.size(0) // 2:]

        return fake, real

    # Given fake and real image, return the prediction of discriminator
    # for each fake and real image.
    def discriminate(self, input_semantics, fake_image, real_image):
        fake_concat = torch.cat([input_semantics, fake_image], dim=1)
        real_concat = torch.cat([input_semantics, real_image], dim=1)

        # In Batch Normalization, the fake and real images are
        # recommended to be in the same batch to avoid disparate
        # statistics in fake and real images.
        # So both fake and real images are fed to D all at once.
        fake_and_real = torch.cat([fake_concat, real_concat], dim=0)
        seg = None
        if self.opt.D_type== 'structure':
            discriminator_out, seg, cam_logit = self.net['netDs'](fake_and_real)
        elif self.opt.D_type=='color':
            discriminator_out, seg, cam_logit = self.net['netDc'](fake_and_real)
        else:
            raise ValueError('can not found' + self.opt.D_type)

        pred_fake, pred_real = self.divide_pred(discriminator_out)
        fake_cam_logit, real_cam_logit = None, None
        if self.opt.D_cam > 0:
            fake_cam_logit = torch.cat([it[:it.shape[0] // 2] for it in cam_logit], dim=1)
            real_cam_logit = torch.cat([it[it.shape[0] // 2:] for it in cam_logit], dim=1)
        # fake_cam_logit, real_cam_logit = self.divide_pred(cam_logit)

        return pred_fake, pred_real, seg, fake_cam_logit, real_cam_logit

    def generate_fake(self, input_semantics, real_image, ref_semantics=None, ref_image=None, self_ref=None):
        generate_out = {}
        # print(ref_image.max())
        ref_relu1_1, ref_relu2_1, ref_relu3_1, ref_relu4_1, ref_relu5_1 = self.vggnet_fix(ref_image,
                                                                                          ['r12', 'r22', 'r32', 'r42',
                                                                                           'r52'], preprocess=True)

        coor_out = self.net['netCorr'](ref_image, real_image, input_semantics, ref_semantics, alpha=self.alpha)

        generate_out['ref_features'] = [ref_relu1_1, ref_relu2_1, ref_relu3_1, ref_relu4_1, ref_relu5_1]
        generate_out['real_features'] = self.vggnet_fix(real_image, ['r12', 'r22', 'r32', 'r42', 'r52'],
                                                        preprocess=True)

        if self.opt.CBN_intype == 'mask':
            CBN_in = input_semantics
        elif self.opt.CBN_intype == 'warp':
            CBN_in = coor_out['warp_out']
        elif self.opt.CBN_intype == 'warp_mask':
            CBN_in = torch.cat((coor_out['warp_out'], input_semantics), dim=1)

        generate_out['fake_image'] = self.net['netG'](input_semantics, warp_out=CBN_in)

        generate_out = {**generate_out, **coor_out}
        return generate_out

    def compute_generator_loss(self, input_label, input_semantics, real_image, ref_label=None, ref_semantics=None,
                               ref_image=None, self_ref=None, similar_ref=None):
        G_losses = {}
        generate_out = self.generate_fake(
            input_semantics, real_image, ref_semantics=ref_semantics, ref_image=ref_image, self_ref=self_ref)

        if 'loss_novgg_featpair' in generate_out and generate_out['loss_novgg_featpair'] is not None:
            G_losses['no_vgg_feat'] = generate_out['loss_novgg_featpair']

        if self.opt.warp_cycle_w > 0:
            if not self.opt.warp_patch:
                ref = F.avg_pool2d(ref_image, self.opt.warp_stride)
            else:
                ref = ref_image

            G_losses['G_warp_cycle'] = F.l1_loss(generate_out['warp_cycle'], ref) * self.opt.warp_cycle_w
            if self.opt.two_cycle:
                real = F.avg_pool2d(real_image, self.opt.warp_stride)
                G_losses['G_warp_cycle'] += F.l1_loss(generate_out['warp_i2r2i'], real) * self.opt.warp_cycle_w

        if self.opt.warp_self_w > 0:
            # real = F.avg_pool2d(real_image, self.opt.warp_stride)
            # warp = F.avg_pool2d(generate_out['warp_out'], self.opt.warp_stride)
            sample_weights = (self_ref[:, 0, 0, 0] / (sum(self_ref[:, 0, 0, 0]) + 1e-5)).unsqueeze(-1).unsqueeze(
                -1).unsqueeze(-1)
            G_losses['G_warp_self'] = torch.mean(
                F.l1_loss(generate_out['warp_out'], real_image, reduce=False) * sample_weights) * self.opt.warp_self_w

        self.opt.D_type = 'structure'
        pred_fake, pred_real, seg, fake_cam_logit, real_cam_logit = self.discriminate(
            input_semantics, generate_out['fake_image'], real_image)
        self.opt.D_type = 'color'
        pred_fake_c, pred_real_c, seg_c, fake_cam_logit_c, real_cam_logit_c = self.discriminate(similar_ref, generate_out['fake_image'], ref_image)
        for i in range(len(pred_fake)):
            num_intermediate_outputs = len(pred_fake[i])
            for j in range(num_intermediate_outputs):
                pred_fake[i][j] = pred_fake[i][j] + pred_fake_c[i][j]
                pred_real[i][j] = pred_real[i][j] + pred_real[i][j]
        if self.opt.D_cam > 0:
            fake_cam_logit.add(fake_cam_logit_c)
            real_cam_logit.add(real_cam_logit_c)
        G_losses['GAN'] = self.criterionGAN(pred_fake, True,
                                            for_discriminator=False) * self.opt.weight_gan

        if not self.opt.no_ganFeat_loss:
            num_D = len(pred_fake)
            GAN_Feat_loss = self.FloatTensor(1).fill_(0)
            for i in range(num_D):  # for each discriminator
                # last output is the final prediction, so we exclude it
                num_intermediate_outputs = len(pred_fake[i]) - 1
                for j in range(num_intermediate_outputs):  # for each layer output
                    unweighted_loss = self.criterionFeat(
                        pred_fake[i][j], pred_real[i][j].detach())
                    GAN_Feat_loss += unweighted_loss * self.opt.lambda_feat / num_D
            G_losses['GAN_Feat'] = GAN_Feat_loss

        fake_features = self.vggnet_fix(generate_out['fake_image'], ['r12', 'r22', 'r32', 'r42', 'r52'],
                                        preprocess=True)
        sample_weights = (self_ref[:, 0, 0, 0] / (sum(self_ref[:, 0, 0, 0]) + 1e-5)).unsqueeze(-1).unsqueeze(
            -1).unsqueeze(-1)
        weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]
        loss = 0
        for i in range(len(generate_out['real_features'])):
            loss += weights[i] * util.weighted_l1_loss(fake_features[i], generate_out['real_features'][i].detach(),
                                                       sample_weights)
        G_losses['fm'] = loss * self.opt.lambda_vgg * self.opt.fm_ratio

        feat_loss = util.mse_loss(fake_features[self.perceptual_layer],
                                  generate_out['real_features'][self.perceptual_layer].detach())
        G_losses['perc'] = feat_loss * self.opt.weight_perceptual

        G_losses['contextual'] = self.get_ctx_loss(fake_features,
                                                   generate_out['ref_features']) * self.opt.lambda_vgg * self.opt.ctx_w

        if self.opt.warp_mask_losstype != 'none':
            ref_label = F.interpolate(ref_label.float(), scale_factor=0.25, mode='nearest').long().squeeze(1)
            gt_label = F.interpolate(input_label.float(), scale_factor=0.25, mode='nearest').long().squeeze(1)
            weights = []
            for i in range(ref_label.shape[0]):
                ref_label_uniq = torch.unique(ref_label[i])
                gt_label_uniq = torch.unique(gt_label[i])
                zero_label = [it for it in gt_label_uniq if it not in ref_label_uniq]
                weight = torch.ones_like(gt_label[i]).float()
                for j in zero_label:
                    weight[gt_label[i] == j] = 0
                weight[gt_label[i] == 0] = 0  # no loss from unknown class
                weights.append(weight.unsqueeze(0))
            weights = torch.cat(weights, dim=0)
            # G_losses['mask'] = (F.cross_entropy(warp_mask, gt_label, reduce =False) * weights.float()).sum() / (weights.sum() + 1e-5) * self.opt.weight_mask
            G_losses['mask'] = (F.nll_loss(torch.log(generate_out['warp_mask'] + 1e-10), gt_label,
                                           reduce=False) * weights).sum() / (
                                       weights.sum() + 1e-5) * self.opt.weight_mask

        # self.fake_image = fake_image
        return G_losses, generate_out

    def preprocess_input(self, data):
        input_semantics = data['label'].clone().cuda().float()
        ref_semantics = data['label_ref'].clone().cuda().float()
        data['label'] = data['label'].cuda()
        data['image'] = data['image'].cuda()
        data['ref'] = data['ref'].cuda()
        data['label_ref'] = data['label_ref'].cuda()
        data['self_ref'] = data['self_ref'].cuda()
        data['similar_ref'] = data['similar_ref'].cuda()
        return data['label'], input_semantics, data['image'], data['self_ref'], data['ref'], data[
            'label_ref'], ref_semantics, data['similar_ref']

    def create_optimizers(self, opt):
        G_params, D_params = list(), list()
        G_params += [{'params': self.net['netG'].parameters(), 'lr': opt.lr * 0.5}]
        G_params += [{'params': self.net['netCorr'].parameters(), 'lr': opt.lr * 0.5}]
        if opt.phase == 'train':
            D_params += list(self.net['netDs'].parameters())
            D_params += list(self.net['netDc'].parameters())
            if opt.weight_domainC > 0 and opt.domain_rela:
                D_params += list(self.net['netDomainClassifier'].parameters())

        if opt.no_TTUR:
            beta1, beta2 = opt.beta1, opt.beta2
            G_lr, D_lr = opt.lr, opt.lr
        else:
            beta1, beta2 = 0, 0.9
            G_lr, D_lr = opt.lr / 2, opt.lr * 2

        optimizer_G = torch.optim.Adam(G_params, lr=G_lr, betas=(beta1, beta2), eps=1e-3)
        optimizer_D = torch.optim.Adam(D_params, lr=D_lr, betas=(beta1, beta2))

        return optimizer_G, optimizer_D

    def initialize_networks(self):
        net = {}
        G = SPADEGenerator(self.opt)
        G.initNetwork(self.opt)
        self.opt.D_type = 'structure'
        Ds = MultiscaleDiscriminator(self.opt)
        Ds.initNetwork(self.opt)
        self.opt.D_type = 'color'
        Dc = MultiscaleDiscriminator(self.opt)
        Dc.initNetwork(self.opt)
        Corr = NoVGGCorrespondence(self.opt)
        Corr.initNetwork(self.opt)
        netDomainClassifier = DomainClassifier(self.opt)
        netDomainClassifier.initNetwork(self.opt)
        net['netG'] = G
        net['netDs'] = Ds if self.opt.phase == 'train' else None
        net['netDc'] = Dc if self.opt.phase == 'train' else None
        net['netCorr'] = Corr
        net[
            'netDomainClassifier'] = netDomainClassifier if self.opt.weight_domainC > 0 and self.opt.domain_rela else None

        # load model
        if self.opt.phase != 'train' or self.opt.continue_train:
            net['netG'] = util.load_network(net['netG'], 'G', self.opt.which_epoch, self.opt)
            if self.opt.phase == 'train':
                net['netDs'] = util.load_network(net['netDs'], 'Ds', self.opt.which_epoch, self.opt)
                net['netDc'] = util.load_network(net['netDc'], 'Dc', self.opt.which_epoch, self.opt)
            net['netCorr'] = util.load_network(net['netCorr'], 'Corr', self.opt.which_epoch, self.opt)
            if self.opt.weight_domainC > 0 and self.opt.domain_rela:
                net['netDomainClassifier'] = util.load_network(net['netDomainClassifier'], 'DomainClassifier',
                                                               self.opt.which_epoch, self.opt)
            if (not self.opt.phase == 'train') and self.opt.use_ema:
                net['netG'] = util.load_network(net['netG'], 'G_ema', self.opt.which_epoch, self.opt)
                net['netCorr'] = util.load_network(net['netCorr'], 'netCorr_ema', self.opt.which_epoch, self.opt)
        return net

    def get_ctx_loss(self, source, target):
        contextual_style5_1 = torch.mean(self.contextual_forward_loss(source[-1], target[-1].detach())) * 8
        contextual_style4_1 = torch.mean(self.contextual_forward_loss(source[-2], target[-2].detach())) * 4
        contextual_style3_1 = torch.mean(self.contextual_forward_loss(F.avg_pool2d(source[-3], 2), F.avg_pool2d(target[-3].detach(), 2))) * 2
        if self.opt.use_22ctx:
            contextual_style2_1 = torch.mean(self.contextual_forward_loss(F.avg_pool2d(source[-4], 4), F.avg_pool2d(target[-4].detach(), 4))) * 1
            return contextual_style5_1 + contextual_style4_1 + contextual_style3_1 + contextual_style2_1
        return contextual_style5_1 + contextual_style4_1 + contextual_style3_1

    def compute_D_seg_loss(self, out, gt):
        fake_seg, real_seg = self.divide_pred([out])
        fake_seg_loss = F.cross_entropy(fake_seg[0][0], gt)
        real_seg_loss = F.cross_entropy(real_seg[0][0], gt)

        down_gt = F.interpolate(gt.unsqueeze(1).float(), scale_factor=0.5, mode='nearest').squeeze().long()
        fake_seg_loss_down = F.cross_entropy(fake_seg[0][1], down_gt)
        real_seg_loss_down = F.cross_entropy(real_seg[0][1], down_gt)

        seg_loss = fake_seg_loss + real_seg_loss + fake_seg_loss_down + real_seg_loss_down
        return seg_loss
if __name__ == '__main__':
    from options.train_options import TrainOptions

    trainoptions = TrainOptions()
    opt = trainoptions.parse()
    bk = Backbone(opt)
